Import all essentials
from global_names import (A2D2_PATH, sensor_p, abs_)
import pickle, random, cv2, os, json, string, imagesize, copy
import numpy as np
from tqdm import tqdm
from matplotlib import pyplot as plt
import matplotlib.patches as patches
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
%matplotlib inline
with open("bm_ds.pkl", "rb") as f:
bm_ds = pickle.load(f)
with open(os.path.join(A2D2_PATH, "camera_lidar_semantic/class_list.json"), 'rb') as f:
class_list = json.load(f)
f = lambda x: int(x, 16)
rgb_int = lambda rgb_list: tuple([i for i in map(f, (rgb_list[1:3], rgb_list[3:5], rgb_list[5:7]))])
class_from_rgb = lambda rgb_list: list(class_list.values())[[list((int(i[1:3], 16), int(i[3:5], 16), int(i[5:7], 16))) for i in class_list.keys()].index(list(rgb_list))]
class_list_simplified = {k: (v[:-2] if v[-1] in string.digits else v) for k,v in class_list.items()}
To convert segmentation maps into 2D bboxes use ./get_bboxes_a2d2.py script (uses 16 threads!)
def run_iter(r_id):
lab_p = abs_(sensor_p(r_id, "label"))
mask_arr = cv2.cvtColor(cv2.imread(lab_p),cv2.COLOR_BGR2RGB)
image_bboxes_dict = {i:[] for i in np.unique(list(class_list_simplified.values()))}
for rgb_list, obj_class in class_list_simplified.items():
binary_mask = np.where(np.prod(np.where(mask_arr == rgb_int(rgb_list), 255, 0), axis=-1) == 0, 0, 255).astype(np.uint8)
contours, _ = cv2.findContours(binary_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# xmin, xmax, ymin, ymax
SIZE_LIMIT = 20
image_bboxes_dict[obj_class] += [i for i in [[c[:,:,0].min(), c[:,:,0].max(), c[:,:,1].min(), c[:,:,1].max()] for i, c in enumerate(contours)] if i[1] - i[0] > SIZE_LIMIT and i[3] - i[2] > SIZE_LIMIT]
all_bboxes[r_id] = image_bboxes_dict
all_ids = []
all_bboxes = {}
for i in bm_ds.values():
all_ids += list(i)
# with ThreadPoolExecutor(16) as executor:
# results = executor.map(run_iter, all_ids)
with open("all_bboxes_small.pkl", "rb") as write_file:
all_bboxes = pickle.load(write_file)
Create whitelist
filter_classes = lambda raw_labels, white_list: {k:v for k,v in raw_labels.items() if k in white_list}
white_list = """Bicycle
Car
Dashed line
Ego car
Irrelevant signs
Obstacles / trash
Pedestrian
Poles
Small vehicles
Solid line
Speed bumper
Tractor
Traffic guide obj.
Traffic sign
Traffic signal
Truck
Utility vehicle
Zebra crossing"""
white_list = white_list.split("\n")
print(white_list)
['Bicycle', 'Car', 'Dashed line', 'Ego car', 'Irrelevant signs', 'Obstacles / trash', 'Pedestrian', 'Poles', 'Small vehicles', 'Solid line', 'Speed bumper', 'Tractor', 'Traffic guide obj.', 'Traffic sign', 'Traffic signal', 'Truck', 'Utility vehicle', 'Zebra crossing']
path = os.path.join("..", "adod", "a2d2_od_dataset")
train = os.path.join("images", "train")
val = os.path.join("images", "val")
test = os.path.join("images", "test")
names = white_list
nc = len(white_list)
dataset_yaml_content = f"""\
path: {path}
train: {train}
val: {val}
test: {test}
nc: {nc}
names: {names}
"""
print(dataset_yaml_content)
path: ../adod/a2d2_od_dataset train: images/train val: images/val test: images/test nc: 18 names: ['Bicycle', 'Car', 'Dashed line', 'Ego car', 'Irrelevant signs', 'Obstacles / trash', 'Pedestrian', 'Poles', 'Small vehicles', 'Solid line', 'Speed bumper', 'Tractor', 'Traffic guide obj.', 'Traffic sign', 'Traffic signal', 'Truck', 'Utility vehicle', 'Zebra crossing']
if not os.path.exists(path):
os.mkdir(path)
for p in [os.path.join(path, i) for i in (train, val, test)]:
if not os.path.exists(p):
os.makedirs(p, exist_ok=False)
with open("../yolov5/a2d2_od.yaml", "w") as f:
f.write(dataset_yaml_content)
for p in [os.path.join(path, i) for i in (train, val, test)]:
p = p.replace("images", "labels")
if not os.path.exists(p):
os.makedirs(p, exist_ok=False)
id_ = random.choice(list(all_bboxes.keys()))
presented_classes = [k for k,v in all_bboxes[id_].items() if len(v) > 0]
img_p = abs_(sensor_p(id_, "camera"))
image_arr = cv2.cvtColor(cv2.imread(img_p),cv2.COLOR_BGR2RGB)
fig, axes = plt.subplots(len(white_list),1, figsize=(10,6*len(white_list)))
for i, class_ in enumerate(white_list):
ax = axes[i]
ax.imshow(image_arr)
for r in [i for i in all_bboxes[id_][class_]]:
rect = patches.Rectangle((r[0], r[2]), r[1]-r[0], r[3]-r[2], linewidth=3, edgecolor='r', facecolor='none')
# Add the patch to the Axes
ax.add_patch(rect)
ax.set_title(class_)
plt.tight_layout()